Skip to content

Metric refactor#524

Merged
nstarman merged 15 commits into
GalacticDynamics:mainfrom
nstarman:metric-refactor
May 18, 2026
Merged

Metric refactor#524
nstarman merged 15 commits into
GalacticDynamics:mainfrom
nstarman:metric-refactor

Conversation

@nstarman
Copy link
Copy Markdown
Contributor

@nstarman nstarman commented May 17, 2026

Metric Refactor Summary

What changed

1. Renamed types for mathematical clarity

Old name New name Why
EuclideanMetric FlatMetric "Flat" describes the geometry, not just the space
HyperSphericalMetric RoundMetric "Round" is the standard term for the sphere's metric
InducedMetric PullbackMetric "Pullback" is the correct mathematical operation
CartesianProductMetric ProductMetric Shorter; "Cartesian product" was redundant
AbstractMetric AbstractMetricField Clarifies that it's a field (assigns a form to each tangent space), not a single matrix
AbstractDiagonalMetric AbstractDiagonalMetricField Consistent with above
QuantityMatrix QMatrix Shorter; no semantic change

2. Three-layer metric architecture replaces method-on-object

Before:

metric.metric_matrix(chart, at=at)  # method on the metric object

After:

cxmapi.metric_matrix(manifold, point, chart)  # standalone dispatch function

The three layers are:

  1. AbstractMetricField — describes the kind of geometry (flat, round, Lorentzian). Does not compute anything.
  2. metric_matrix(manifold, point, chart) — dispatch function that computes and returns the matrix, registered per (manifold, chart) pair.
  3. AbstractMetricMatrix (DiagonalMetric / DenseMetric) — typed result carrying the matrix and its sparsity structure.

A companion function metric_representation(manifold, chart) queries which type will be returned without allocating.

3. Typed metric matrix results

The old implementation returned a raw QuantityMatrix (always dense). The new implementation returns:

  • DiagonalMetric — stores only the 1-D diagonal; O(n) matmul instead of O(n²)
  • DenseMetric — full matrix; used for non-orthogonal charts and pullback metrics

change_basis automatically selects the efficient path based on the result type.

4. RoundMetric becomes an Equinox module

Unlike the other metrics (frozen static dataclasses), RoundMetric can carry a dynamic radius leaf, enabling JIT and grad through radius parameters — important for astronomy use cases.

5. Functional-only API for scale_factors and angle_between

Previously these were also available as convenience methods on metric and manifold objects
(metric.scale_factors(...), M.angle_between(...)). They are now purely functional, which
simplifies the class hierarchy and removes the need to call through the object.

6. Pre-defined manifold singletons

# Before
cxm.EuclideanManifold(3), cxm.HyperSphericalManifold()
# After
cxm.R3, cxm.S2  # (also R1, R2, etc.)

Why it's better

Concern Old design New design
Separation of concerns Metric field also computed the matrix via metric_matrix method Field describes geometry; computation is in external dispatch rules
Sparsity exploitation Always returned a full QuantityMatrix Returns DiagonalMetric for orthogonal charts — O(n) ops vs O(n²)
Extensibility Adding a new chart required subclassing the metric Register a new metric_matrix dispatch rule without touching existing types
Query without allocation No way to ask "what type will I get?" metric_representation(M, chart) answers without computing
Naming accuracy EuclideanMetric implied the manifold; InducedMetric hid the operation FlatMetric, RoundMetric, PullbackMetric use standard math vocabulary
Dynamic parameters All metrics were frozen/static RoundMetric can carry a differentiable radius

nstarman added 11 commits May 17, 2026 18:21
… package

QuantityMatrix is renamed to QMatrix for brevity. The monolithic quantity_matrix.py module is supplemented by a new quantity_matrix/ package that splits the implementation into focused submodules (_quantity_matrix, _units_matrix, _det, _inv, _register_primitives, _utils). det and inv operations are added to the public surface. The .py file and the package coexist; Python resolves to the package at import time.
…rix types, and dispatch API

Adds a new coordinax._src.metric package with four modules: field.py (AbstractMetricField, AbstractDiagonalMetricField, RoundMetric), matrix.py (AbstractMetricMatrix, DiagonalMetric, DenseMetric), api.py (metric_matrix, metric_representation fallback dispatches), and __init__.py. AbstractMetric/AbstractDiagonalMetric in base/metric.py are renamed to AbstractMetricField/AbstractDiagonalMetricField and stripped of their metric_matrix method — that computation is now a standalone dispatch function.
…e_factors, and rename metric classes

Each geometry (Euclidean, HyperSpherical, Minkowski, Product, Embedded, Custom) gains a register_metric.py module that registers metric_matrix and metric_representation dispatches using the new standalone dispatch API. Metric class names are updated: EuclideanMetric → FlatMetric, HyperSphericalMetric → RoundMetric, CartesianProductMetric → ProductMetric, InducedMetric → PullbackMetric. Spherical and Minkowski geometries receive new scale_factors.py files. The metric_matrix method is removed from each metric class in favour of the dispatch function.
…ic symbols

The coordinax.api.manifolds dispatch API gains metric_matrix and metric_representation. The coordinax.manifolds re-export module is updated to expose the renamed metric classes (FlatMetric, RoundMetric, ProductMetric, PullbackMetric) and new matrix types (AbstractMetricMatrix, DiagonalMetric, DenseMetric). AbstractManifold is updated to use AbstractMetricField. The scale_factors dispatch file gains metric-aware overloads. Public internal.py and main.py re-exports are aligned with the new names.
…, representations, transforms, vectors, distances, and frames

All references to QuantityMatrix (now QMatrix), EuclideanMetric (now FlatMetric), HyperSphericalMetric (now RoundMetric), CartesianProductMetric (now ProductMetric), InducedMetric (now PullbackMetric), AbstractMetric (now AbstractMetricField), and AbstractDiagonalMetric (now AbstractDiagonalMetricField) are updated throughout the core modules and their tests. Docstring Examples section headers that merely introduce the code block are also removed for cleaner inline doctests.
…ng cleanup

Remove Examples section headers from docstrings in distance_modulus, frame_transforms, parallax, register_constructors, and register_converters. Update tests to use new metric and frame API conventions.
…ld API

Remove unused register_frames dispatches and update tests to match the current frame and manifold API conventions.
…anifold API

Update strategy tests and chart/vector/distance tests to reflect the current API conventions introduced by the metric refactor.
…ocstring cleanup

Replace QuantityMatrix with QMatrix in qmatrix.py and remove Examples section headers from angles, distances, frames, ptmap, vec_constructors, and vec_converters. Tests updated to match.
spec.md is updated to reflect the new metric module design (AbstractMetricField, metric_matrix dispatch, DiagonalMetric/DenseMetric matrix types). API and guide pages for manifolds, charts, and representations are updated to use the new class names. README and tutorials are adjusted for minor API changes.
Add jax.extend.core and jax.tree_util import abbreviations to pyproject.toml. Add ARG002 (unused method argument) to per-file test ignores to allow pytest fixture parameters that are injected but not referenced in the body.
@nstarman nstarman added this to the v0.24.0 milestone May 17, 2026
@github-actions github-actions Bot added 🔧 Add / update configuration Add or update configuration files. 📝 Add / update documentation Add or update documentation. ✅ Add / update / pass tests Add, update, or pass tests. ♻️ Refactor code Refactor code. ✨ Introduce new features Introduce new features. labels May 17, 2026
nstarman added 3 commits May 17, 2026 18:48
Cast u.ustrip() results to Array/jax.Array so .ndim/.shape resolve; cast q.unit to u.AbstractUnit; fix ty: ignore placement for J.T attribute access; add ty: ignore[possibly-missing-submodule] on jax.core references.
Replace float literals (1.0, 2.0, etc.) with integer literals where values are whole numbers, and normalize rtol=0.0 to rtol=0 in assert_allclose calls.
…ity_matrix tests

Replace float literals like 1.0, 2.0 with int equivalents where the values are whole numbers. Preserve floats where needed for JAX grad/jvp tests, dtype-sensitive operations, and mixed-unit conversion tests that produce float results.
@nstarman nstarman marked this pull request as ready for review May 17, 2026 23:35
Copilot AI review requested due to automatic review settings May 17, 2026 23:35
Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR is a significant refactor of the metric system in coordinax. It renames types for mathematical clarity (e.g., EuclideanMetricFlatMetric, HyperSphericalMetricRoundMetric, QuantityMatrixQMatrix), restructures the metric API around a three-layer architecture (AbstractMetricField + dispatched metric_matrix function + typed AbstractMetricMatrix results), and introduces pre-defined manifold singletons (cxm.R3, cxm.S2). The diff hunks shown are test-file updates that propagate these renames and adjust assertions to the new typed-result API.

Changes:

  • Renames QuantityMatrixQMatrix, EuclideanMetricFlatMetric, HyperSphericalMetricRoundMetric, InducedMetricPullbackMetric across tests.
  • Updates the scale-factors test to use the new metric_matrix dispatch and DiagonalMetric result type instead of the removed metric.metric_matrix(...) method.
  • Replaces cxm.HyperSphericalManifold()/cxm.EuclideanManifold(3) usages with singletons cxm.S2/cxm.R3, and changes various float literals (1.0) to int literals (1) in test data.

Reviewed changes

Copilot reviewed 174 out of 174 changed files in this pull request and generated 1 comment.

Show a summary per file
File Description
tests/usage/charts/test_jacobian.py Renames QuantityMatrixQMatrix in imports/asserts; changes vmap test inputs to int literals.
tests/unit/vectors/test_tangent.py Switches batched tangent component data to int literals.
tests/unit/transforms/conftest.py Renames QuantityMatrixQMatrix in fixture; uses int literals for fixture data and expected tuples.
tests/unit/representations/test_tangent_map.py Reformats dicts and converts test data to int literals.
tests/unit/manifolds/test_scale_factors_dispatch.py Adopts new metric_matrix dispatch + DiagonalMetric, renames metric/manifold classes, and renames two test functions to embed QMatrix in their names.
tests/unit/manifolds/test_custom.py Replaces EuclideanMetric with FlatMetric in CustomManifold construction.
tests/unit/manifolds/test_angle_between_dispatch.py Replaces EuclideanMetric/HyperSphericalMetric with FlatMetric/RoundMetric.
tests/unit/distances/test_distance.py Converts float literals to int literals in hypothesis strategies, parametrize cases, and assertions.
tests/unit/charts/test_utils.py Converts float literals in test arrays to int literals.

Comment thread tests/unit/manifolds/test_scale_factors_dispatch.py
…egies and ptmap

Replace float literals like 8.0, -8.0, -3.0, 3.0, -5.0, 5.0, 10.0, -10.0 with int equivalents where the values are whole numbers and dtype safety permits.

Signed-off-by: nstarman <nstarman@users.noreply.github.com>
@codecov
Copy link
Copy Markdown

codecov Bot commented May 18, 2026

Codecov Report

❌ Patch coverage is 89.80632% with 100 lines in your changes missing coverage. Please review.
✅ Project coverage is 92.91%. Comparing base (9480275) to head (aed8dea).

Files with missing lines Patch % Lines
...c/internal/quantity_matrix/_register_primitives.py 83.67% 16 Missing ⚠️
...nax/_src/internal/quantity_matrix/_units_matrix.py 87.69% 16 Missing ⚠️
src/coordinax/representations/_src/basis_change.py 69.76% 13 Missing ⚠️
src/coordinax/_src/manifolds/scale_factors.py 76.59% 11 Missing ⚠️
src/coordinax/_src/minkowski/register_metric.py 70.96% 9 Missing ⚠️
src/coordinax/_src/euclidean/register_metric.py 95.72% 5 Missing ⚠️
.../_src/internal/quantity_matrix/_quantity_matrix.py 92.42% 5 Missing ⚠️
src/coordinax/_src/metric/api.py 61.53% 5 Missing ⚠️
...rc/coordinax/_src/internal/quantity_matrix/_det.py 92.00% 4 Missing ⚠️
...rc/coordinax/_src/internal/quantity_matrix/_inv.py 90.69% 4 Missing ⚠️
... and 8 more
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #524      +/-   ##
==========================================
- Coverage   93.24%   92.91%   -0.33%     
==========================================
  Files         232      249      +17     
  Lines        7386     7877     +491     
==========================================
+ Hits         6887     7319     +432     
- Misses        499      558      +59     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@nstarman nstarman merged commit e3b9782 into GalacticDynamics:main May 18, 2026
16 of 18 checks passed
@nstarman nstarman deleted the metric-refactor branch May 18, 2026 02:48
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

🔧 Add / update configuration Add or update configuration files. 📝 Add / update documentation Add or update documentation. ✅ Add / update / pass tests Add, update, or pass tests. ✨ Introduce new features Introduce new features. ♻️ Refactor code Refactor code.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants